package com.ics.atable.chat.service;

import jakarta.annotation.PreDestroy;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * embedding service
 *
 * @author haohongbin
 */
@Service
public class EmbeddingService {

    private static final int DEFAULT_CHUNK_SIZE = 2048;
    private final EmbeddingModel embeddingModel;
    private final ExecutorService virtualThreadExecutor;


    public EmbeddingService(EmbeddingModel embeddingModel) {

        this.embeddingModel = embeddingModel;
        // 使用Java 21虚拟线程创建执行器
        this.virtualThreadExecutor = Executors.newVirtualThreadPerTaskExecutor();
    }

    public EmbeddingResponse embedding(List<String> messages) {
        return this.embeddingModel.embedForResponse(messages);
    }


    /**
     * 同步 embedding 方法（分块处理）
     * @param messages 消息队列
     * @return embedding结果
     */
    public EmbeddingResponse embeddingByChunk(List<String> messages) {
        List<Embedding> embeddings = new ArrayList<>();
        EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();

        for (int i = 0; i < messages.size(); i += DEFAULT_CHUNK_SIZE) {
            List<String> subList = messages.subList(i, Math.min(i + DEFAULT_CHUNK_SIZE, messages.size()));
            EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(subList);
            embeddings.addAll(embeddingResponse.getResults());
            metadata = embeddingResponse.getMetadata();
        }

        return new EmbeddingResponse(embeddings, metadata);
    }

    /**
     * 分块并行处理大量消息
     *
     * @param messages 需要处理的消息列表
     * @param chunkSize 每块大小
     * @return CompletableFuture<EmbeddingResponse> 异步处理结果
     */
    public CompletableFuture<EmbeddingResponse> embeddingChunkedAsync(List<String> messages, int chunkSize) {
        if (messages.isEmpty()) {
            return CompletableFuture.completedFuture(new EmbeddingResponse(List.of()));
        }

        // 将消息分块
        List<List<String>> chunks = new ArrayList<>();
        for (int i = 0; i < messages.size(); i += chunkSize) {
            int end = Math.min(i + chunkSize, messages.size());
            chunks.add(messages.subList(i, end));
        }

        // 并行处理每个块
        List<CompletableFuture<EmbeddingResponse>> futures = chunks.stream()
                .map(chunk -> CompletableFuture.supplyAsync(
                        () -> this.embeddingModel.embedForResponse(chunk),
                        virtualThreadExecutor))
                .toList();

        // 合并结果
        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .thenApply(v -> {
                    List<EmbeddingResponse> responses = futures.stream()
                            .map(CompletableFuture::join)
                            .toList();

                    // 合并所有响应
                    return mergeEmbeddingResponses(responses);
                });
    }

    /**
     * 使用默认块大小的分块并行处理
     * @param messages 需要处理的消息列表
     * @return CompletableFuture<EmbeddingResponse> 异步处理结果
     */
    public CompletableFuture<EmbeddingResponse> embeddingChunkedAsync(List<String> messages) {
        return embeddingChunkedAsync(messages, DEFAULT_CHUNK_SIZE);
    }

    // 合并多个EmbeddingResponse的辅助方法
    private EmbeddingResponse mergeEmbeddingResponses(List<EmbeddingResponse> responses) {
        // 合并所有Embedding
        List<Embedding> mergedEmbeddings = responses.stream()
                .flatMap(response -> response.getResults().stream())
                .toList();

        EmbeddingResponseMetadata metadata = responses.getFirst().getMetadata();

        return new EmbeddingResponse(mergedEmbeddings, metadata);
    }

    /**
     * 并行处理embedding的接口，使用虚拟线程+CompletableFuture
     *
     * @param messages 需要处理的消息列表
     * @return CompletableFuture<EmbeddingResponse> 异步处理结果
     */
    public CompletableFuture<EmbeddingResponse> embeddingAsync(List<String> messages) {
        return CompletableFuture.supplyAsync(() -> this.embeddingModel.embedForResponse(messages),
                virtualThreadExecutor);
    }

    /**
     * 批量并行处理多个消息列表
     *
     * @param messageLists 多个消息列表
     * @return CompletableFuture<List<EmbeddingResponse>> 异步处理结果列表
     */
    public CompletableFuture<List<EmbeddingResponse>> embeddingBatchAsync(List<List<String>> messageLists) {
        // 为每个消息列表创建一个异步任务
        List<CompletableFuture<EmbeddingResponse>> futures = messageLists.stream()
                .map(messages -> CompletableFuture.supplyAsync(() -> this.embeddingModel.embedForResponse(messages),
                        virtualThreadExecutor))
                .toList();

        // 等待所有任务完成并收集结果
        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .thenApply(v -> futures.stream()
                        .map(CompletableFuture::join)
                        .toList());
    }

    /**
     * 关闭虚拟线程执行器
     */
    public void shutdown() {
        if (virtualThreadExecutor != null && !virtualThreadExecutor.isShutdown()) {
            virtualThreadExecutor.shutdown();
        }
    }

    @PreDestroy
    public void cleanup() {
        shutdown();
    }
}
